import argparse
import atexit
import logging
import select
import socket
import struct
import subprocess
import sys
import time
from socketserver import StreamRequestHandler, ThreadingTCPServer
from threading import Thread

SOCKS_VERSION = 5


class Socket5Proxy(StreamRequestHandler):
    def handle(self):
        logging.info(f"new request from {self.client_address[0]}:{self.client_address[1]}")

        self._check_request()

        try:
            address, port = self._get_target_addr_port()
            remote = self._remote_connect(address, port)
        except Exception as err:
            logging.error(err)
            self._reply_faild(1)
        else:
            self._exchange_data(self.connection, remote)
        finally:
            self.server.close_request(self.request)

    def _reply_faild(self, error_number):
        reply = struct.pack("!BBBBIH", SOCKS_VERSION, error_number, 0, 1, 0, 0)
        self.connection.sendall(reply)

    def _remote_connect(self, address, port):
        remote = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        remote.connect((address, port))
        bind_address = remote.getsockname()
        logging.debug(f"address: {address}, port: {port}, bind_address: {bind_address[0]}:{bind_address[1]}")
        addr = struct.unpack("!I", socket.inet_aton(bind_address[0]))[0]
        port = bind_address[1]
        reply = struct.pack("!BBBBIH", SOCKS_VERSION, 0, 0, 1, addr, port)
        self.connection.sendall(reply)
        return remote

    def _exchange_data(self, client, remote):
        while True:
            logging.debug("block for reading ...")
            rs, _, _ = select.select([client, remote], [], [])
            if client in rs:
                data = client.recv(4096)
                logging.debug(f"read from client and put to remote, size={len(data)}")
                if remote.send(data) <= 0:
                    break
            if remote in rs:
                data = remote.recv(4096)
                logging.debug(f"read from remote and put to client, size={len(data)}")
                if client.send(data) <= 0:
                    break

    def _check_request(self):
        header = self.connection.recv(2)
        VER, NMETHODS = struct.unpack("!BB", header)
        if VER != SOCKS_VERSION:
            raise Exception(f"version of socket:{VER} is unsupported, expected is {SOCKS_VERSION}")

        # 仅支持无密码连接
        methods = {ord(self.connection.recv(1)) for _ in range(NMETHODS)}
        if 0 not in set(methods):
            raise Exception(f"only support non-passwd-logging")
        else:
            self.connection.sendall(struct.pack("!BB", SOCKS_VERSION, 0))

    def _get_target_addr_port(self):
        _, cmd, _, address_type = struct.unpack(
            "!BBBB", self.connection.recv(4))

        if cmd != 1:
            raise Exception(f"unsupported CMD:{cmd}, expected is 1")

        if address_type == 1:    # ip地址, 对应: export http_proxy="socks5://host:port"
            address = socket.inet_ntoa(self.connection.recv(4))
        elif address_type == 3:    # 域名, 对应: export http_proxy="socks5h://host:port"
            domain_length = ord(self.connection.recv(1))
            address = self.connection.recv(domain_length).decode()
        else:
            raise Exception(f"unsupported address_type:{address_type}, expected in [1, 3]")

        port = struct.unpack('!H', self.connection.recv(2))[0]
        return address, port


def _get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--ssh_cmd', type=str,
                        default='user@host -p port')
    parser.add_argument('--socket5_port', type=int, default=3089)
    parser.add_argument('--logger_level', type=logging.getLevelName, default='INFO')
    args = parser.parse_args()
    return args


def _start_ssh_tunnel(args):
    ssh_full_cmd = f"ssh -N -R 127.0.0.1:{args.socket5_port}:127.0.0.1:{args.socket5_port}  {args.ssh_cmd}"
    logging.info(f"ssh cmd: [ {ssh_full_cmd} ]")
    ssh_process = subprocess.Popen(args=ssh_full_cmd, shell=True, close_fds=False,
                                   stdin=sys.stdin, stdout=sys.stdout, stderr=sys.stdout)
    atexit.register(ssh_process.terminate)

    def _flush_stdout():
        for _ in range(120):
            sys.stdout.flush()
            time.sleep(1)
        logging.info("finish flushing stdout os ssh process")

    Thread(target=_flush_stdout, daemon=True).start()


def _welcome(args):
    logging.info(f"start socket5 proxy 127.0.0.1:{args.socket5_port}")
    logging.info(f"command of setting proxy in linux:")
    print("#######################猴哥制造，放心良药#######################")
    print()
    print(f"export http_proxy=\"socks5h://127.0.0.1:{args.socket5_port}\"")
    print(f"export https_proxy=\"socks5h://127.0.0.1:{args.socket5_port}\"")
    print()
    print("##############################################################")


if __name__ == '__main__':
    args = _get_args()

    logging.basicConfig(level=args.logger_level,
                        format=">> [%(levelname)s] [%(asctime)s]-%(filename)s-line:%(lineno)d: %(message)s")

    try:

        _welcome(args)

        _start_ssh_tunnel(args)

        ThreadingTCPServer(('127.0.0.1', args.socket5_port), Socket5Proxy).serve_forever()

    except BaseException as e:
        logging.error(e)
        input("input anything to exit...")
